package com.qfmy.config;


import com.alibaba.fastjson.JSONObject;
import com.qfmy.service.assistant.ChatService;
import common.login.LoginUser;

import common.login.LoginUserHolder;
import common.utils.JwtUtil;
import io.jsonwebtoken.Claims;
import lombok.extern.slf4j.Slf4j;

import org.jetbrains.annotations.NotNull;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import org.springframework.util.MultiValueMap;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Flux;


import java.io.IOException;
import java.net.URLDecoder;
import java.nio.charset.StandardCharsets;

/**
 * WebSocket配置类
 */
@Slf4j
@Configuration
@EnableWebSocket
@SuppressWarnings("all")
public class WebSocketConfig implements WebSocketConfigurer {

    @Autowired
    private ChatService chatService; // 假设你有一个ChatService处理聊天逻辑

    /**
     * 注册WebSocket处理器
     * @param registry
     */
    @Override
    public void registerWebSocketHandlers(@NotNull WebSocketHandlerRegistry registry) {
        registry.addHandler(myHandler(), "/app/chat/stream/websocket", "/app/chat/image")
                .setAllowedOrigins("*");
    }

    /**
     * 创建WebSocket处理器
     * @return
     */
    @Bean
    public WebSocketHandler myHandler() {
        return new TextWebSocketHandler() {
            @Override
            protected void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
                try {
                    String payload = message.getPayload();
                    JSONObject json = JSONObject.parseObject(payload);
                    String type = json.getString("type");

                    // 1. 处理认证消息
                    if ("auth".equals(type)) {
                        String token = json.getString("token");
                        if (token == null) {
                            session.close(CloseStatus.POLICY_VIOLATION);
                            return;
                        }

                        // 验证Token
                        Claims claims = JwtUtil.parseToken(token);
                        if (claims == null) {
                            session.close(CloseStatus.POLICY_VIOLATION);
                            return;
                        }

                        // 提取用户信息并保存到Session
                        String phone = (String) claims.get("phone");
                        String username = (String) claims.get("username");
                        Long userId = Long.valueOf((Integer) claims.get("userId"));
                        LoginUser loginUser = new LoginUser(phone, userId, username);
                        //将用户信息保存在线程里面
                        LoginUserHolder.setLoginUser(loginUser);
                    }

                    // 2. 从URL中提取业务参数
                    String query = session.getUri().getQuery();
                    if (query == null || query.isEmpty()) {
                        session.close(CloseStatus.POLICY_VIOLATION);
                        return;
                    }

                    // 解析URL参数
                    MultiValueMap<String, String> params = UriComponentsBuilder
                            .fromUriString(session.getUri().toString())
                            .build()
                            .getQueryParams();

                    // 提取业务参数
                    String prompt = params.getFirst("prompt");
                    if (prompt == null) {
                        session.close(CloseStatus.POLICY_VIOLATION);
                        return;
                    }
                    prompt = URLDecoder.decode(prompt, StandardCharsets.UTF_8);

                    boolean isThink = Boolean.parseBoolean(params.getFirst("isThink"));
                    boolean search = Boolean.parseBoolean(params.getFirst("search"));

                    String[] image = new String[0];
                    String imageParam = params.getFirst("image");
                    if (imageParam != null) {
                        imageParam = URLDecoder.decode(imageParam, StandardCharsets.UTF_8);
                        image = imageParam.split(",");
                    }

                    // 4. 调用业务逻辑
                    Flux<String> responseFlux;
                    if (image == null || image.length == 0) {
                        responseFlux = chatService.text(prompt, isThink, search);
                    } else {
                        responseFlux = chatService.image(image, prompt);
                    }

                    // 5. 订阅流并发送消息
                    responseFlux.subscribe(
                            msg -> {
                                try {
                                    session.sendMessage(new TextMessage(msg));
                                } catch (IOException e) {
                                    log.error("发送消息失败", e);
                                    throw new RuntimeException(e);
                                }
                            },
                            error -> {
                                try {
                                    session.sendMessage(new TextMessage("Error: " + error.getMessage()));
                                    session.close();
                                } catch (IOException e) {
                                    log.error("发送消息失败", e);
                                    throw new RuntimeException(e);
                                }
                            },
                            () -> {
                                try {
                                    session.close();
                                } catch (IOException e) {
                                    log.error("关闭WebSocket会话失败", e);
                                    throw new RuntimeException(e);
                                }
                            }
                    );

                } catch (Exception e) {
                    log.error("处理消息失败", e);
                    session.close();
                }
            }
        };
    }
}